# Author : ZZH
# Date : 2025/5/26
import cv2
from ultralytics import YOLO
from pathlib import Path
import argparse
from utils import custom_plot
from utils import process_frame
from utils import (setup_logging, load_yaml_config, merge_configs,
                log_execution_time, rename_log_file, YOLO_SERVICE_DIR, CONFIGS_DIR)
import logging


def args_parser():
    parser = argparse.ArgumentParser(description="基于YOLO的安全帽推理脚本")

    parser.add_argument("--weights", type = str, default = "train6-20250523-175204-yolo11m-best.pt", help = "模型权重路径")
    parser.add_argument("--source", type = str, default = "0", help = "输入源（图像/文件夹/视频/摄像头ID，如‘0’）")
    parser.add_argument("--imgsz", type = int, default = 640, help = "输入图像尺寸")
    parser.add_argument("--conf", type = float, default = 0.25, help = "置信度阈值")
    parser.add_argument("--iou", type = float, default = 0.45, help = "IOU阈值")
    parser.add_argument("--save", type = bool, default = True, help = "保存预测图像")
    parser.add_argument("--save_dir", type = str, default = "runs/val", help = "保存预测图像的目录")
    parser.add_argument("--save_txt", type = bool, default = True, help = "保存预测结果为TXT")
    parser.add_argument("--save_conf", type = bool, default = True, help = "在TXT中包含置信度值")
    parser.add_argument("--save_frames", type = bool, default = True, help = "保存摄像头/视频每帧图像")
    parser.add_argument("--save_crop", type = bool, default = True, help = "保存检测框裁剪图像")
    parser.add_argument("--show", type = bool, default = True, help = "显示结果")
    parser.add_argument("--display-size",
                        type=str,
                        default="720p",
                        choices=["360p", "720p", "1280p", "2K", "4K"],
                        help="显示窗口大小")
    parser.add_argument("--beautify", type=bool, default=True, help="美化结果")
    parser.add_argument("--use-chinese",
                        type=bool,
                        default=True,
                        help="使用中文标签映射 (默认启用)")
    # 美化参数组
    parser.add_argument("--font-size",
                        type=int,
                        default=20,
                        help="基准字体大小（以720p为基准自动缩放）")
    parser.add_argument("--line-width",
                        type=int,
                        default=4,
                        help="边界框线宽（以720p为基准自动缩放）")
    parser.add_argument("--label-padding-x",
                        type=int,
                        default=30,
                        help="美化标签水平内边距（像素单位）")
    parser.add_argument("--label-padding-y",
                        type=int,
                        default=18,
                        help="美化标签垂直内边距（像素单位）")
    parser.add_argument("--radius",
                        type=int,
                        default=8,
                        help="美化圆角半径（像素单位）")
    parser.add_argument("--text-color",
                        type=str,
                        default="black",
                        help="文本颜色名称")
    parser.add_argument("--log-encoding",
                        type=str,
                        default="utf-8-sig",
                        choices=["utf-8", "utf-8-sig", "gbk"],
                        help="日志文件编码格式")
    parser.add_argument("--use_yaml",
                        type=bool,
                        default=True,
                        help="使用YAML配置文件（默认启用）")
    parser.add_argument("--log_level", type=str, default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"], help="日志记录级别")
    return parser.parse_args()


def main():
    """
    主函数，执行推理
    :return: None
    """
    args = args_parser()
    # 设置日志级别
    log_level = getattr(logging, args.log_level.upper(), logging.INFO)
    # 获取模型标识
    model_name = Path(args.weights).stem
    # 初始化日志系统
    logger = setup_logging(
        base_path=YOLO_SERVICE_DIR,
        log_type="infer",
        model_name=model_name,
        encoding=args.log_encoding,
        log_level=log_level,
        temp_log=True
    )
    # 记录推理开始
    logger.info("==== YOLOv8 安全帽检测推理开始 =====")
    # 加载配置文件
    try:
        yaml_config = {}
        if args.use_yaml:
            yaml_config = load_yaml_config(
                config_path=CONFIGS_DIR / "infer.yaml",
                config_type="infer"
            )
        # 合并配置参数
        yolo_args, project_args = merge_configs(args, yaml_config, mode="infer")

        # 分辨率的映射
        resolution_map = {
            "360p": (640, 360),
            "720p": (1280, 720),
            "1280p": (1920, 1080),  # 注意：常规1080p应为1920x1080
            "2K": (2560, 1440),
            "4K": (3840, 2160)
        }
        # 获取显示尺寸
        display_width, display_height = resolution_map[args.display_size]

        # 检测结果美化参数配置
        beautify_params = {
            "use_chinese_mapping": project_args.use_chinese,
            "beautify": project_args.beautify,
            "font_size": project_args.font_size,
            "line_width": project_args.line_width,
            "label_padding": (project_args.label_padding_x,project_args.label_padding_y),
            "radius": project_args.radius,
            "text_color": (0, 0, 0)
        }

        # 模型路径标准化
        model_path = Path(project_args.weights)
        if not model_path.is_absolute():
            model_path = YOLO_SERVICE_DIR / "models" / "checkpoints" / project_args.weights
        # 输入源验证
        source = project_args.source

        if not source.isdigit():  # 摄像头ID为数字时跳过路径检查
            source_path = Path(source)
            if not source_path.exists():
                logger.error(f"输入源不存在：{source_path}")
                raise FileNotFoundError(f"输入源不存在：{source_path}")
            source = str(source_path)
        # 模型存在性检查
        if not model_path.exists():
            logger.error(f"模型文件不存在：{model_path}")
            raise FileNotFoundError(f"模型文件不存在：{project_args.weights}")
        # 加载模型
        logger.info(f"加载模型：{model_path}")
        model = YOLO(str(model_path))

        window_name = "YOLO Safety Helmet Detection"
        cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)

        # 摄像头推理
        if source.isdigit() or source.endswith((".mp4", "*.avi", "*.mov")):
            # 初始化视频写入
            cap = cv2.VideoCapture(int(source) if source.isdigit() else source)
            if not cap.isOpened():
                logger.error(f"无法打开{'摄像头' if source.isdigit() else '视频'}: {source}")
                raise RuntimeError(f"无法打开{'摄像头' if source.isdigit() else '视频'}: {source}")
            fps = cap.get(cv2.CAP_PROP_FPS) or 30

            preprocess_times = []
            inference_times = []
            postprocess_times = []
            frame_count = 0
            # 流式推理
            video_writer = None
            frames_dir = None
            save_dir = None
            # 参数设置：复制yolo_args参数，修改为流模式
            yolo_args_dict = vars(yolo_args).copy()
            yolo_args_dict['stream'] = True
            yolo_args_dict['show'] = False
            yolo_args_dict['save'] = False
            # # 设置显示窗口
            # window_name = "YOLO Safety Helmet Detection"
            # cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
            # cv2.resizeWindow(window_name, display_width, display_height)

            # 逐帧推理
            for idx, result in enumerate(model.predict(**yolo_args_dict)):
                # 第一帧初始化保存路径
                if idx == 0:
                    # 初始化结果保存目录
                    save_dir = Path(result.save_dir)
                    logger.info(f"推理结果保存目录: {save_dir}")
                    # 创建帧图像存储目录
                    if project_args.save_frames:
                        frames_dir = save_dir / "0_frames"
                        frames_dir.mkdir(parents=True, exist_ok=True)
                    # 初始化视频写入器
                    if project_args.save:
                        video_path = save_dir / "output.mp4"
                        video_writer = cv2.VideoWriter(
                            str(video_path),
                            cv2.VideoWriter_fourcc(*"mp4v"),
                            fps,
                            (display_width, display_height)
                        )
                # 获取原始帧和标注结果
                frame = result.orig_img
                # 缩放标注帧
                annotated_frame = process_frame(frame, result, args, display_width, display_height, beautify_params)

                # 写入视频文件
                if video_writer:
                    video_writer.write(annotated_frame)
                if frames_dir:
                    frame_path = frames_dir / f"{idx}.jpg"
                    cv2.imwrite(str(frame_path), annotated_frame)
                    logger.debug(f"保存帧图像:{frame_path}")

                # 收集推理时间
                preprocess_times.append(result.speed['preprocess'])
                inference_times.append(result.speed['inference'])
                postprocess_times.append(result.speed['postprocess'])
                frame_count += 1
                # 显示
                cv2.imshow(window_name, annotated_frame)
                key = cv2.waitKey(10) & 0xFF
                if key == ord('q') or key == 27:
                    logger.info("用户按q或Esc退出推理")
                    break

            # 计算总时间和平均时间
            total_preprocess = sum(preprocess_times)
            total_inference = sum(inference_times)
            total_postprocess = sum(postprocess_times)
            total_time = total_preprocess + total_inference + total_postprocess

            avg_preprocess = total_preprocess / frame_count if frame_count > 0 else 0
            avg_inference = total_inference / frame_count if frame_count > 0 else 0
            avg_postprocess = total_postprocess / frame_count if frame_count > 0 else 0
            avg_total = total_time / frame_count if frame_count > 0 else 0

            # 日志记录时间
            logger.info(
                f"视频/摄像头推理总时间：总帧数={frame_count}, "
                f"总预处理={total_preprocess:.2f}ms, "
                f"总推理={total_inference:.2f}ms, "
                f"总后处理={total_postprocess:.2f}ms, "
                f"总耗时={total_time:.2f}ms (约{total_time / 1000:.2f}秒)"
            )

            logger.info(
                f"单帧平均时间："
                f"预处理={avg_preprocess:.2f}ms, "
                f"推理={avg_inference:.2f}ms, "
                f"后处理={avg_postprocess:.2f}ms, "
                f"总计={avg_total:.2f}ms"
            )
            # 释放资源
            cap.release()  # 释放视频捕获对象
            if video_writer:
                video_writer.release()  # 释放视频写入对象
            cv2.destroyAllWindows()     # 关闭所有OpenCV窗口
            # 记录推理完成日志
            logger.info(f"{'摄像头' if source.isdigit() else '视频'}推理完成，结果已保存至: {save_dir or '未保存'}")
            # 重命名文件
            if save_dir:
                rename_log_file(logger, save_dir, model_name, project_args.log_encoding)
        else:
            # 非流式推理（图像/文件夹处理）
            yolo_args_dict = vars(yolo_args).copy()
            yolo_args_dict.pop('stream', None)
            yolo_args_dict['show'] = False
            yolo_args_dict['save'] = False
            # 执行模型推理
            results = model.predict(**yolo_args_dict)
            # 初始化保存目录
            save_dir = Path(results[0].save_dir)
            logger.info(f"推理结果保存目录: {save_dir}")

            # 处理每张推理结果
            for result in results:
                # 生成标注图像
                display_height,display_width = result.orig_shape[:2]
                annotated_frame = process_frame(
                    result.orig_img, result, project_args, display_width, display_height, beautify_params
                )
                # 保存标注结果
                if project_args.save:
                    output_path = save_dir / Path(result.path).name
                    cv2.imwrite(str(output_path), annotated_frame)
                    logger.debug(f"保存图像: {output_path}")

                # 记录处理时间
                logger.info(
                    f"图像 {Path(result.path).name}: "
                    f"预处理={result.speed['preprocess']:.2f}ms, "
                    f"推理={result.speed['inference']:.2f}ms, "
                    f"后处理={result.speed['postprocess']:.2f}ms"
                )

                # 显示交互
                cv2.imshow(window_name, annotated_frame)
                key = cv2.waitKey(0) & 0xFF
                if key == ord('q') or key == 27:
                    logger.info("用户按 q 或 Esc 退出推理")
                    break

            # 清理资源
            cv2.destroyAllWindows()
            logger.info(f"推理完成，结果已保存至: {save_dir}")
            rename_log_file(logger, save_dir, model_name, project_args.log_encoding)
    except Exception as e:
        logger.error(f"参数解析错误: {e}")
        raise
    finally:
        logger.info("===== YOLOv8 安全帽检测推理结束 =====")


if __name__ == "__main__":
    main()
